package kvraft

import (
	"../labgob"
	"../labrpc"
	"../raft"
	"sync"
	"sync/atomic"
)

type Op struct {
	// Your definitions here.
	// Field names must start with capital letters,
	// otherwise RPC will break.
	Command string
	Key     string
	Value   string
}

type KVServer struct {
	mu      sync.Mutex
	me      int
	rf      *raft.Raft
	applyCh chan raft.ApplyMsg
	dead    int32 // set by Kill()

	maxraftstate int // snapshot if log grows this big

	// Your definitions here.
	storage map[string]string
	dm      map[int]chan bool
	xt      map[int64]int

	done chan bool
	e    chan interface{}
}

type ApplyMsgEvent struct {
	applyMsg raft.ApplyMsg
}

func (kv *KVServer) Get(args *GetArgs, reply *GetReply) {
	// Your code here.
	command := Op{
		Command: GET,
		Key:     args.Key,
	}
	index, _, isLeader := kv.rf.Start(command)
	if !isLeader {
		reply.Err = ErrWrongLeader
		return
	}

	done := make(chan bool)
	kv.mu.Lock()
	kv.dm[index] = done
	kv.mu.Unlock()
	<-done

	kv.mu.Lock()
	defer kv.mu.Unlock()

	value, ok := kv.storage[args.Key]
	if ok {
		reply.Value = value
	} else {
		reply.Err = ErrNoKey
	}
}

func (kv *KVServer) PutAppend(args *PutAppendArgs, reply *PutAppendReply) {
	kv.mu.Lock()
	if index, ok := kv.xt[args.Id]; ok {
		done, wait := kv.dm[index]
		kv.mu.Unlock()
		if wait {
			<-done
		}
		return
	}
	kv.mu.Unlock()
	command := Op{
		// Your code here.
		Command: args.Op,
		Key:     args.Key,
		Value:   args.Value,
	}
	index, _, isLeader := kv.rf.Start(command)
	if !isLeader {
		reply.Err = ErrWrongLeader
		return
	}

	done := make(chan bool)
	kv.mu.Lock()
	kv.xt[args.Id] = index
	kv.dm[index] = done
	kv.mu.Unlock()
	<-done

	kv.mu.Lock()
	defer kv.mu.Unlock()
	if kv.doPutAppend(command) {
		reply.Err = OK
	} else {
		reply.Err = ErrNoKey
	}
}

func (kv *KVServer) doPutAppend(op Op) bool {
	switch op.Command {
	case PUT:
		kv.storage[op.Key] = op.Value
	case APPEND:
		if _, ok := kv.storage[op.Key]; !ok {
			return false
		}
		DPrintf("%d->Before Append:%s", kv.me, kv.storage[op.Key])
		kv.storage[op.Key] += op.Value
		DPrintf("%d->After Append:%s", kv.me, kv.storage[op.Key])
	}
	return true
}

//
// the tester calls Kill() when a KVServer instance won't
// be needed again. for your convenience, we supply
// code to set rf.dead (without needing a lock),
// and a killed() method to test rf.dead in
// long-running loops. you can also add your own
// code to Kill(). you're not required to do anything
// about this, but it may be convenient (for example)
// to suppress debug output from a Kill()ed instance.
//
func (kv *KVServer) Kill() {
	atomic.StoreInt32(&kv.dead, 1)
	kv.rf.Kill()
	// Your code here, if desired.
	close(kv.done)
}

func (kv *KVServer) killed() bool {
	z := atomic.LoadInt32(&kv.dead)
	return z == 1
}

func (kv *KVServer) loop() {
	DPrintf("%d->server start", kv.me)
	for !kv.killed() {
		select {
		case <-kv.done:
			return

		case applyMsg := <-kv.applyCh:
			kv.processApplyMsgEvent(applyMsg)
		}
	}
	DPrintf("%d->server killed", kv.me)
}

func (kv *KVServer) processApplyMsgEvent(applyMsg raft.ApplyMsg) {
	kv.mu.Lock()
	defer kv.mu.Unlock()

	index := applyMsg.CommandIndex
	if ch, ok := kv.dm[index]; ok {
		DPrintf("%d->ApplyMsg:%+v", kv.me, applyMsg)
		close(ch)
		delete(kv.dm, index)
	} else {
		op := applyMsg.Command.(Op)
		kv.doPutAppend(op)
	}
}

//
// servers[] contains the ports of the set of
// servers that will cooperate via Raft to
// form the fault-tolerant key/value service.
// me is the index of the current server in servers[].
// the k/v server should storage snapshots through the underlying Raft
// implementation, which should call persister.SaveStateAndSnapshot() to
// atomically save the Raft state along with the snapshot.
// the k/v server should snapshot when Raft's saved state exceeds maxraftstate bytes,
// in order to allow Raft to garbage-collect its log. if maxraftstate is -1,
// you don't need to snapshot.
// StartKVServer() must return quickly, so it should start goroutines
// for any long-running work.
//
func StartKVServer(servers []*labrpc.ClientEnd, me int, persister *raft.Persister, maxraftstate int) *KVServer {
	// call labgob.Register on structures you want
	// Go's RPC library to marshall/unmarshall.
	labgob.Register(Op{})

	kv := new(KVServer)
	kv.me = me
	kv.maxraftstate = maxraftstate

	// You may need initialization code here.

	kv.applyCh = make(chan raft.ApplyMsg)
	kv.rf = raft.Make(servers, me, persister, kv.applyCh)

	// You may need initialization code here.
	kv.storage = make(map[string]string)
	kv.dm = make(map[int]chan bool)
	kv.xt = make(map[int64]int)
	kv.done = make(chan bool)
	kv.e = make(chan interface{})

	go kv.loop()

	return kv
}
